import numpy as np

def predict_action(state: np.ndarray):
    """
    策略3：VWAP + MFI + ROC + EMA 组合动量策略
    输入:
        state: np.ndarray of shape (b, 32, 20)
    输出:
        action: np.ndarray of shape (b, 1)，取值为 [-1, 0, 1]
    """
    b, l, c = state.shape
    assert l == 32 and c == 20

    IDX = {
        "close": 4,
        "ema_12": 8,
        "ema_26": 9,
        "roc": 13,
        "mfi": 16,
        "vwap": 19
    }

    close = state[:, -1, IDX["close"]]
    vwap = state[:, -1, IDX["vwap"]]

    mfi = state[:, -1, IDX["mfi"]]
    mfi_prev = state[:, -2, IDX["mfi"]]

    roc = state[:, -1, IDX["roc"]]
    roc_prev = state[:, -2, IDX["roc"]]

    ema_12 = state[:, -1, IDX["ema_12"]]
    ema_26 = state[:, -1, IDX["ema_26"]]

    # ==== 多头信号 ====
    long_cond = (
        (close < vwap) &
        (mfi > 50) & (mfi > mfi_prev) &
        (roc > 0) & (roc > roc_prev) &
        (ema_12 > ema_26)
    )

    # ==== 空头信号 ====
    short_cond = (
        (close > vwap) &
        (mfi < 50) & (mfi < mfi_prev) &
        (roc < 0) & (roc < roc_prev) &
        (ema_12 < ema_26)
    )

    action = np.zeros(b, dtype=np.int32)
    action[long_cond] = 1
    action[short_cond] = -1

    return action.reshape(-1, 1)
